#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
bots.py —— 单进程并发 N 个“机器人客户端”，按你自定义协议压服务端
协议: $ + type(2) + ip(4) + len(4) + data + #
模式:
  - img (默认): 生成JPEG → zlib压缩 → Base64，再以 IMG_SEND(0) 发送
  - text      : 文本 -> zlib压缩，再以 TEXT_SEND(4) 发送（与你的 TEXT_RECV 解压对齐）
用法示例:
  python3 bots.py --server 127.0.0.1 --port 1234 --room 1732 --n 100 --mode text --pps 5
  python3 bots.py --server 127.0.0.1 --port 1234 --room 9700 --n 100 --mode img --fps 15 --w 320 --h 240 --jpeg 60
"""
import asyncio, struct, time, os, zlib, base64, argparse, random, sys
from io import BytesIO

# 可选: 生成测试帧（没有 Pillow 也能跑 text 模式；img 模式建议安装）
try:
    from PIL import Image, ImageDraw
except Exception:
    Image = None

# ---- 默认参数（可用命令行覆盖） ----
DEFAULT_SERVER = "127.0.0.1"
DEFAULT_PORT   = 9000
DEFAULT_ROOM   = 1234
DEFAULT_N      = 100
DEFAULT_MODE   = "img"    # img | text
DEFAULT_FPS    = 15.0     # img 模式帧率
DEFAULT_PPS    = 5        # text 模式每秒包数
DEFAULT_W      = 320
DEFAULT_H      = 240
DEFAULT_JPEG_Q = 60
LOCAL_IP_BE    = 0        # IP字段（大端）; 服务端不校验可填0

# ---- 按你的 netheader.h 的真实枚举值 ----
IMG_SEND      = 0
IMG_RECV      = 1
AUDIO_SEND    = 2
AUDIO_RECV    = 3
TEXT_SEND     = 4
TEXT_RECV     = 5
CREATE_MEETING= 6
EXIT_MEETING  = 7
JOIN_MEETING  = 8
CLOSE_CAMERA  = 9
# -----------------------------------

# IMG：你的接收端代码是 Base64 + qUncompress，所以这里发送 Base64(zlib(JPEG))
# TEXT：你的接收端对 TEXT_RECV 直接 qUncompress(cc)（无Base64），所以这里仅 zlib
def pack(msg_type: int, ip_be: int, payload: bytes) -> bytes:
    return (b'$' +
            struct.pack('>H', msg_type) +
            struct.pack('>I', ip_be) +
            struct.pack('>I', len(payload)) +
            payload + b'#')

def make_jpeg(w:int, h:int, ts:float, q:int) -> bytes:
    if Image is None:
        return (f"{ts:.3f}".encode()).ljust(1024, b'\0')  # 占位
    img = Image.new('RGB', (w, h), (0, 0, 0))
    d = ImageDraw.Draw(img)
    d.rectangle([0, 0, w, 36], fill=(255, 255, 255))
    d.text((8, 8), f"{ts:.3f}", fill=(0, 0, 0))
    buf = BytesIO(); img.save(buf, format='JPEG', quality=q)
    return buf.getvalue()

def encode_img(jpg: bytes) -> bytes:
    return base64.b64encode(zlib.compress(jpg))

def encode_text(s: str) -> bytes:
    return zlib.compress(s.encode('utf-8'))

async def bot(idx:int, server:str, port:int, room:int,
              mode:str, fps:float, pps:int, w:int, h:int, jpeg_q:int):
    try:
        reader, writer = await asyncio.open_connection(server, port)
    except Exception as e:
        print(f"[bot#{idx}] connect fail: {e}")
        return

    # JOIN_MEETING：data=room(uint32, 大端)
    try:
        writer.write(pack(JOIN_MEETING, LOCAL_IP_BE, struct.pack('>I', room)))
        await writer.drain()
    except Exception as e:
        print(f"[bot#{idx}] join send fail: {e}")
        with contextlib.suppress(Exception):
            writer.close(); await writer.wait_closed()
        return

    sent = 0
    t_win = time.time()
    period = 1.0 / (fps if mode == "img" else max(pps, 1))
    await asyncio.sleep(random.uniform(0, period))  # 轻微错峰

    try:
        while True:
            t_loop = time.time()
            if mode == "img":
                jpg = make_jpeg(w, h, t_loop, jpeg_q)
                payload = encode_img(jpg)
                pkt = pack(IMG_SEND, LOCAL_IP_BE, payload)
            else:
                payload = encode_text(f"bot#{idx} {t_loop:.3f}")
                pkt = pack(TEXT_SEND, LOCAL_IP_BE, payload)

            writer.write(pkt)
            await writer.drain()
            sent += 1

            # 每5秒打印一次速率
            now = time.time()
            if now - t_win >= 5.0:
                rate = sent / (now - t_win)
                print(f"[bot#{idx}] {mode} sent={sent}  rate≈{rate:.1f}/s")
                sent = 0; t_win = now

            # 控频
            elapsed = time.time() - t_loop
            await asyncio.sleep(max(0.0, period - elapsed))
    except asyncio.CancelledError:
        pass
    except Exception as e:
        print(f"[bot#{idx}] loop error: {e}")
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass

async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", default=os.getenv("CM_SERVER_IP", DEFAULT_SERVER))
    ap.add_argument("--port",   type=int, default=int(os.getenv("CM_SERVER_PORT", str(DEFAULT_PORT))))
    ap.add_argument("--room",   type=int, default=int(os.getenv("CM_ROOM", str(DEFAULT_ROOM))))
    ap.add_argument("--n",      type=int, default=int(os.getenv("CM_N", str(DEFAULT_N))))
    ap.add_argument("--mode",   choices=["img","text"], default=os.getenv("CM_MODE", DEFAULT_MODE))
    ap.add_argument("--fps",    type=float, default=float(os.getenv("CM_FPS", str(DEFAULT_FPS))))
    ap.add_argument("--pps",    type=int,   default=int(os.getenv("CM_PPS", str(DEFAULT_PPS))))
    ap.add_argument("--w",      type=int,   default=int(os.getenv("CM_W",   str(DEFAULT_W))))
    ap.add_argument("--h",      type=int,   default=int(os.getenv("CM_H",   str(DEFAULT_H))))
    ap.add_argument("--jpeg",   type=int,   default=int(os.getenv("CM_JPEG",str(DEFAULT_JPEG_Q))))
    args = ap.parse_args()

    if args.mode == "img" and Image is None:
        print("[warn] 未安装 Pillow，将发送占位字节；建议 `pip3 install pillow` 以发送真实JPEG帧。")

    tasks = []
    for i in range(args.n):
        tasks.append(asyncio.create_task(
            bot(i, args.server, args.port, args.room,
                args.mode, args.fps, args.pps, args.w, args.h, args.jpeg)
        ))
        await asyncio.sleep(0.01)  # 建连错峰

    print(f"launched {args.n} bots → {args.server}:{args.port} room={args.room} mode={args.mode}")
    try:
        await asyncio.gather(*tasks)
    except KeyboardInterrupt:
        for t in tasks: t.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)

if __name__ == "__main__":
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())
